#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
屏幕锁定检测和解锁功能模块
用于处理云主机锁屏情况
"""

import time
import win32gui
import win32con
import win32api
import win32clipboard
from ctypes import windll, wintypes, byref, c_int
import ctypes
from OutPut.outPut import op


class ScreenLockManager:
    """屏幕锁定管理器"""
    
    def __init__(self):
        self.user32 = windll.user32
        self.kernel32 = windll.kernel32
        
    def is_screen_locked(self):
        """检测屏幕是否被锁定
        
        Returns:
            bool: True表示屏幕已锁定，False表示未锁定
        """
        try:
            # 方法1: 检查是否有活动的屏幕保护程序
            screensaver_active = self.user32.SystemParametersInfoW(
                win32con.SPI_GETSCREENSAVERRUNNING, 0, None, 0
            )
            
            # 方法2: 检查桌面是否可访问
            desktop_accessible = self._is_desktop_accessible()
            
            # 方法3: 检查是否存在锁屏窗口
            lock_screen_exists = self._check_lock_screen_window()
            
            # 方法4: 检查工作站是否被锁定
            workstation_locked = self._is_workstation_locked()
            
            is_locked = screensaver_active or not desktop_accessible or lock_screen_exists or workstation_locked
            
            if is_locked:
                op(f"[!] 检测到屏幕已锁定 - 屏保:{screensaver_active}, 桌面不可访问:{not desktop_accessible}, 锁屏窗口:{lock_screen_exists}, 工作站锁定:{workstation_locked}")
            else:
                op(f"[*] 屏幕未锁定")
                
            return is_locked
            
        except Exception as e:
            op(f"[!] 检测屏幕锁定状态时发生异常: {str(e)}")
            return False
    
    def _is_desktop_accessible(self):
        """检查桌面是否可访问"""
        try:
            # 尝试获取桌面窗口
            desktop_hwnd = self.user32.GetDesktopWindow()
            if not desktop_hwnd:
                return False
                
            # 尝试获取前台窗口
            foreground_hwnd = self.user32.GetForegroundWindow()
            if not foreground_hwnd:
                return False
                
            return True
        except:
            return False
    
    def _check_lock_screen_window(self):
        """检查是否存在锁屏相关窗口"""
        try:
            # 常见的锁屏窗口类名
            lock_screen_classes = [
                'Windows.UI.Core.CoreWindow',  # Windows 10/11 锁屏
                'LockScreenControllerProxyWindow',
                'LockScreenController',
                'Windows.UI.Core.CoreWindow',
                'ApplicationFrameWindow'
            ]
            
            for class_name in lock_screen_classes:
                hwnd = win32gui.FindWindow(class_name, None)
                if hwnd:
                    window_text = win32gui.GetWindowText(hwnd)
                    if any(keyword in window_text.lower() for keyword in ['lock', '锁定', 'signin', '登录']):
                        return True
            
            return False
        except:
            return False
    
    def _is_workstation_locked(self):
        """检查工作站是否被锁定"""
        try:
            # 使用WTSQuerySessionInformation检查会话状态
            import ctypes
            from ctypes import wintypes
            
            # 定义常量
            WTS_CURRENT_SESSION = -1
            WTSConnectState = 8
            
            # 加载wtsapi32.dll
            wtsapi32 = ctypes.windll.wtsapi32
            
            # 查询会话状态
            session_info = wintypes.DWORD()
            bytes_returned = wintypes.DWORD()
            
            result = wtsapi32.WTSQuerySessionInformationW(
                wintypes.HANDLE(0),  # WTS_CURRENT_SERVER_HANDLE
                WTS_CURRENT_SESSION,
                WTSConnectState,
                ctypes.byref(ctypes.pointer(session_info)),
                ctypes.byref(bytes_returned)
            )
            
            if result:
                # 0 = Active, 1 = Connected, 2 = ConnectQuery, 3 = Shadow, 4 = Disconnected, 5 = Idle, 6 = Listen, 7 = Reset, 8 = Down, 9 = Init
                # 4 (Disconnected) 通常表示锁定状态
                return session_info.contents.value == 4
            
            return False
        except:
            return False
    
    def unlock_screen(self, password=None):
        """解锁屏幕
        
        Args:
            password (str, optional): 解锁密码
            
        Returns:
            bool: 解锁是否成功
        """
        try:
            op(f"[*] 开始尝试解锁屏幕...")
            
            # 如果屏幕没有锁定，直接返回成功
            if not self.is_screen_locked():
                op(f"[*] 屏幕未锁定，无需解锁")
                return True
            
            # 步骤1: 唤醒屏幕
            self._wake_screen()
            time.sleep(1)
            
            # 步骤2: 发送按键唤醒
            self._send_wake_keys()
            time.sleep(2)
            
            # 步骤3: 如果提供了密码，输入密码
            if password:
                success = self._input_password(password)
                if not success:
                    op(f"[!] 密码输入失败")
                    return False
            else:
                op(f"[*] 未提供密码，尝试无密码解锁")
                # 尝试按回车键（适用于无密码的情况）
                win32api.keybd_event(win32con.VK_RETURN, 0, 0, 0)
                win32api.keybd_event(win32con.VK_RETURN, 0, win32con.KEYEVENTF_KEYUP, 0)
            
            # 步骤4: 等待并检查解锁结果
            time.sleep(3)
            
            # 验证解锁是否成功
            unlock_success = not self.is_screen_locked()
            
            if unlock_success:
                op(f"[*] 屏幕解锁成功!")
            else:
                op(f"[!] 屏幕解锁失败")
                
            return unlock_success
            
        except Exception as e:
            op(f"[!] 解锁屏幕时发生异常: {str(e)}")
            return False
    
    def _wake_screen(self):
        """唤醒屏幕"""
        try:
            # 方法1: 移动鼠标
            current_pos = win32gui.GetCursorPos()
            win32api.SetCursorPos((current_pos[0] + 1, current_pos[1]))
            time.sleep(0.1)
            win32api.SetCursorPos(current_pos)
            
            # 方法2: 发送鼠标移动消息
            self.user32.mouse_event(win32con.MOUSEEVENTF_MOVE, 1, 1, 0, 0)
            time.sleep(0.1)
            self.user32.mouse_event(win32con.MOUSEEVENTF_MOVE, -1, -1, 0, 0)
            
        except Exception as e:
            op(f"[!] 唤醒屏幕时发生异常: {str(e)}")
    
    def _send_wake_keys(self):
        """发送唤醒按键"""
        try:
            # 按空格键唤醒
            win32api.keybd_event(win32con.VK_SPACE, 0, 0, 0)
            win32api.keybd_event(win32con.VK_SPACE, 0, win32con.KEYEVENTF_KEYUP, 0)
            time.sleep(0.5)
            
            # 按Ctrl键（某些系统需要）
            win32api.keybd_event(win32con.VK_CONTROL, 0, 0, 0)
            win32api.keybd_event(win32con.VK_CONTROL, 0, win32con.KEYEVENTF_KEYUP, 0)
            
        except Exception as e:
            op(f"[!] 发送唤醒按键时发生异常: {str(e)}")
    
    def _input_password(self, password):
        """输入密码
        
        Args:
            password (str): 要输入的密码
            
        Returns:
            bool: 输入是否成功
        """
        try:
            op(f"[*] 正在输入解锁密码...")
            
            # 等待密码输入框出现
            time.sleep(2)
            
            # 清空可能存在的内容
            win32api.keybd_event(win32con.VK_CONTROL, 0, 0, 0)
            win32api.keybd_event(ord('A'), 0, 0, 0)
            win32api.keybd_event(ord('A'), 0, win32con.KEYEVENTF_KEYUP, 0)
            win32api.keybd_event(win32con.VK_CONTROL, 0, win32con.KEYEVENTF_KEYUP, 0)
            time.sleep(0.2)
            
            # 输入密码
            for char in password:
                if char.isalnum():
                    # 字母数字直接输入
                    vk_code = ord(char.upper())
                    win32api.keybd_event(vk_code, 0, 0, 0)
                    win32api.keybd_event(vk_code, 0, win32con.KEYEVENTF_KEYUP, 0)
                else:
                    # 特殊字符通过剪贴板输入
                    self._send_char_via_clipboard(char)
                time.sleep(0.1)
            
            # 按回车确认
            time.sleep(0.5)
            win32api.keybd_event(win32con.VK_RETURN, 0, 0, 0)
            win32api.keybd_event(win32con.VK_RETURN, 0, win32con.KEYEVENTF_KEYUP, 0)
            
            return True
            
        except Exception as e:
            op(f"[!] 输入密码时发生异常: {str(e)}")
            return False
    
    def _send_char_via_clipboard(self, char):
        """通过剪贴板发送字符"""
        try:
            # 保存当前剪贴板内容
            old_clipboard = ""
            try:
                win32clipboard.OpenClipboard()
                old_clipboard = win32clipboard.GetClipboardData()
                win32clipboard.CloseClipboard()
            except:
                pass
            
            # 设置字符到剪贴板
            win32clipboard.OpenClipboard()
            win32clipboard.EmptyClipboard()
            win32clipboard.SetClipboardText(char)
            win32clipboard.CloseClipboard()
            time.sleep(0.1)
            
            # 粘贴
            win32api.keybd_event(win32con.VK_CONTROL, 0, 0, 0)
            win32api.keybd_event(ord('V'), 0, 0, 0)
            win32api.keybd_event(ord('V'), 0, win32con.KEYEVENTF_KEYUP, 0)
            win32api.keybd_event(win32con.VK_CONTROL, 0, win32con.KEYEVENTF_KEYUP, 0)
            time.sleep(0.1)
            
            # 恢复剪贴板内容
            if old_clipboard:
                try:
                    win32clipboard.OpenClipboard()
                    win32clipboard.EmptyClipboard()
                    win32clipboard.SetClipboardText(old_clipboard)
                    win32clipboard.CloseClipboard()
                except:
                    pass
                    
        except Exception as e:
            op(f"[!] 通过剪贴板发送字符时发生异常: {str(e)}")
    
    def auto_unlock_with_retry(self, password=None, max_retries=3):
        """带重试的自动解锁
        
        Args:
            password (str, optional): 解锁密码
            max_retries (int): 最大重试次数
            
        Returns:
            bool: 解锁是否成功
        """
        from wechat_config import config
        
        max_retries = config.SCREEN_UNLOCK_RETRIES
        retry_delay = config.SCREEN_UNLOCK_RETRY_DELAY
        
        for attempt in range(max_retries):
            try:
                op(f"[*] 解锁尝试 {attempt + 1}/{max_retries}")
                
                success = self.unlock_screen(password)
                if success:
                    op(f"[*] 第 {attempt + 1} 次尝试解锁成功!")
                    return True
                    
                if attempt < max_retries - 1:
                    op(f"[*] 第 {attempt + 1} 次尝试失败，{retry_delay}秒后重试...")
                    time.sleep(retry_delay)
                    
            except Exception as e:
                op(f"[!] 第 {attempt + 1} 次解锁尝试异常: {str(e)}")
                if attempt < max_retries - 1:
                    time.sleep(retry_delay)
        
        op(f"[!] 所有解锁尝试均失败")
        return False


# 全局实例
screen_lock_manager = ScreenLockManager()
